// Adapted from https://raw.githubusercontent.com/microsoft/monaco-editor/main/src/basic-languages/python/python.ts.
/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

export function registerLmqlLanguage(monaco) {
    const conf = {
        comments: {
            lineComment: '#',
            blockComment: ["'''", "'''"]
        },
        brackets: [
            ['{', '}'],
            ['[', ']'],
            ['(', ')']
        ],
        autoClosingPairs: [
            { open: '{', close: '}' },
            { open: '[', close: ']' },
            { open: '(', close: ')' },
            { open: '"', close: '"', notIn: ['string'] },
            { open: "'", close: "'", notIn: ['string', 'comment'] }
        ],
        surroundingPairs: [
            { open: '{', close: '}' },
            { open: '[', close: ']' },
            { open: '(', close: ')' },
            { open: '"', close: '"' },
            { open: "'", close: "'" }
        ],
        onEnterRules: [
            {
                beforeText: new RegExp(
                    '^\\s*(?:def|class|for|if|elif|else|while|try|with|finally|except|async|match|case).*?:\\s*$'
                ),
                action: { indentAction: monaco.languages.IndentAction.Indent }
            }
        ],
        folding: {
            offSide: true,
            markers: {
                start: new RegExp('^\\s*#region\\b'),
                end: new RegExp('^\\s*#endregion\\b')
            }
        }
    };

    const language = {
        defaultToken: '',
        tokenPostfix: '.python',

        keywords: [
            "BEAM",
            "beam",
            "ARGMAX",
            "argmax",
            "incontext",
            "SAMPLE",
            "BEST_K",
            "best_k",
            "BEAM_VAR",
            "beam_var",
            "VAR",
            "var",
            "sample",
            "FROM",
            "from",
            "WHERE",
            "where",
            "DISTRIBUTION",
            "distribution",
            "context",
        
            // This section is the result of running
            // `import keyword; for k in sorted(keyword.kwlist + keyword.softkwlist): print("  '" + k + "',")`
            // in a Python REPL,
            // though note that the output from Python 3 is not a strict superset of the
            // output from Python 2.
            'False', // promoted to keyword.kwlist in Python 3
            'None', // promoted to keyword.kwlist in Python 3
            'True', // promoted to keyword.kwlist in Python 3
            '_', // new in Python 3.10
            'and',
            'as',
            'assert',
            'async', // new in Python 3
            'await', // new in Python 3
            'break',
            'case', // new in Python 3.10
            'class',
            'continue',
            'def',
            'del',
            'elif',
            'else',
            'except',
            'exec', // Python 2, but not 3.
            'finally',
            'for',
            'from',
            'global',
            'if',
            'import',
            'in',
            'is',
            'lambda',
            'match', // new in Python 3.10
            'nonlocal', // new in Python 3
            'not',
            'or',
            'pass',
            'print', // Python 2, but not 3.
            'raise',
            'return',
            'try',
            'while',
            'with',
            'yield',

            'int',
            'float',
            'long',
            'complex',
            'hex',

            'abs',
            'all',
            'any',
            'apply',
            'basestring',
            'bin',
            'bool',
            'buffer',
            'bytearray',
            'callable',
            'chr',
            'classmethod',
            'cmp',
            'coerce',
            'compile',
            'complex',
            'delattr',
            'dict',
            'dir',
            'divmod',
            'enumerate',
            'eval',
            'execfile',
            'file',
            'filter',
            'format',
            'frozenset',
            'getattr',
            'globals',
            'hasattr',
            'hash',
            'help',
            'id',
            'input',
            'intern',
            'isinstance',
            'issubclass',
            'iter',
            'len',
            'locals',
            'list',
            'map',
            'max',
            'memoryview',
            'min',
            'next',
            'object',
            'oct',
            'open',
            'ord',
            'pow',
            'print',
            'property',
            'reversed',
            'range',
            'raw_input',
            'reduce',
            'reload',
            'repr',
            'reversed',
            'round',
            'self',
            'set',
            'setattr',
            'slice',
            'sorted',
            'staticmethod',
            'str',
            'sum',
            'super',
            'tuple',
            'type',
            'unichr',
            'unicode',
            'vars',
            'xrange',
            'zip',

            '__dict__',
            '__methods__',
            '__members__',
            '__class__',
            '__bases__',
            '__name__',
            '__mro__',
            '__subclasses__',
            '__init__',
            '__import__'
        ],

        brackets: [
            { open: '{', close: '}', token: 'delimiter.curly' },
            { open: '[', close: ']', token: 'delimiter.bracket' },
            { open: '(', close: ')', token: 'delimiter.parenthesis' }
        ],

        tokenizer: {
            root: [
                { include: '@whitespace' },
                { include: '@numbers' },
                { include: '@strings' },

                [/[,:;]/, 'delimiter'],
                [/[{}\[\]()]/, '@brackets'],

                [/@[a-zA-Z_]\w*/, 'tag'],
                [
                    /[a-zA-Z_]\w*/,
                    {
                        cases: {
                            '@keywords': 'keyword',
                            '@default': 'identifier'
                        }
                    }
                ]
            ],

            // Deal with white space, including single and multi-line comments
            whitespace: [
                [/\s+/, 'white'],
                [/(^#.*$)/, 'comment'],
                // use ''' but not '''lmql
                [/'''[^(lmql)]/, 'string', '@endDocString'],
                [/'''lmql/, 'comment'],
                [/'''$/, 'comment'],
                [/"""/, 'string', '@endDblDocString']
            ],
            endDocString: [
                { include: '@queryString' },
                [/[^'\[\{]/, 'string'],
                [/\\'/, 'string'],
                [/'''/, 'string', '@popall'],
                [/'/, 'string']
            ],
            endLmqlDocString: [
                // everything like top-level
                // [/[^(''')]+/, 'root'],
                [/'''/, 'comment', '@popall'],
            ],
            endDblDocString: [  
                { include: '@queryString' },
                [/[^"\[\{]/, 'string'],
                [/\\"/, 'string'],
                [/"""/, 'string', '@popall'],
                [/"/, 'string'],
            ],

            // Recognize hex, negatives, decimals, imaginaries, longs, and scientific notation
            numbers: [
                [/-?0x([abcdef]|[ABCDEF]|\d)+[lL]?/, 'number.hex'],
                [/-?(\d*\.)?\d+([eE][+\-]?\d+)?[jJ]?[lL]?/, 'number']
            ],

            // Recognize strings, including those broken across lines with \ (but not without)
            strings: [
                [/r'[^']*'/, 'string.regex'],
                [/r"[^"]*"/, 'string.regex'],
                [/'$/, 'string.escape', '@popall'],
                [/'/, 'string.escape', '@stringBody'],
                [/"$/, 'string.escape', '@popall'],
                [/"/, 'string.escape', '@dblStringBody']
            ],
            stringBody: [
                { include: '@queryString' },
                
                [/[^\\']+$/, 'string', '@popall'],
                [/\[\[+/, 'string'],
                [/[^\\'\[\{]/, 'string'],
                [/\\./, 'string'],
                [/'/, 'string.escape', '@popall'],
                [/\\$/, 'string']
            ],
            dblStringBody: [
                { include: '@queryString' },
                [/[^\\"]+$/, 'string', '@popall'],
                [/[^\\"\[\{}]/, 'string'],
                [/\\./, 'string'],
                [/"/, 'string.escape', '@popall'],
                [/\\$/, 'string']
            ],
            queryString: [
                [/r'[^']*'/, 'string.regex'],
                [/r"[^"]*"/, 'string.regex'],
                // check for regex
                [/\[\[/, 'string'],
                [/\{\{/, 'string'],
                [/\[/, 'string', '@placeholderVar'],
                [/\{/, 'string', '@templateVar']
            ],
            // var: type or just var
            placeholderVar: [
                [/:/, 'storage', '@pop'],
                [/\]/, 'string', '@pop'],
                [/[^\]:]+/, 'storage']
            ],
            templateVar: [
                [/:[^}]+/, 'tag'],
                [/\}/, 'string', '@pop'],
                [/[^\}]+/, 'variable.parameter']
            ],
        }
    };
    
    // define custom monaco language
    monaco.languages.register({ id: 'lmql' });
    // register language
    monaco.languages.setMonarchTokensProvider('lmql', language);
    // set configuration
    monaco.languages.setLanguageConfiguration('lmql', conf);
}